import random

import torch
import torch.nn as nn

import torchvision.transforms.functional
import matplotlib.pyplot as plt


class torch_nn_layer_usage_class(nn.Module):
    """
    定义几个常用的层，并演示这些层的使用效果
    """

    def __init__(self):
        """
        一般层的定义都是在__init__函数中进行的。
        一般必须先运行 ``super(Name, self).__init__()``
        """
        super(torch_nn_layer_usage_class, self).__init__()
        """
        #-------------------------------卷积层--------------------------------#
        nn.Conv2d(
            in_channels: int,
            out_channels: int,
            kernel_size: _size_2_t,
            stride: _size_2_t = 1,
            padding: _size_2_t = 0,
            dilation: _size_2_t = 1,
            groups: int = 1,
            bias: bool = True,
            padding_mode: str = 'zeros'  # TODO: refine this type
            )
        # in_channels: 输入通道数
        # out_channels: 输出通道数
        # kernel_size: 卷积核尺寸，整数或者元组
        # stride: 卷积操作的步幅，整数或者元组
        # padding: 数据hw方向上填充的层数，整数或者元组，默认填充的是0
        # dilation: 卷积核内部各点的间距，整数或者元组
        # groups: 控制输入和输出之间的连接-
            group=1，输出是所有输入的卷积；
            group=2，此时相当于有并排的两个卷积层，每个卷积层计算输入通道的一半，并且产生的输出是输出通道的一半，随后将这两个输出连接起来。
        # bias: bool类型，为true的话就会对输出添加可学习的偏置量
        # padding_mode: 填充模式，字符串类型
        """
        in_channels = 3
        out_channels = 6
        kernel_size = (5, 3)
        self.conv = nn.Conv2d(in_channels=in_channels,
                              out_channels=out_channels,
                              kernel_size=kernel_size)

        """
        #-------------------------------最大池化层--------------------------------#
        MaxPool2d(
            kernel_size: _size_any_t, 
            stride: Optional[_size_any_t] = None,
            padding: _size_any_t = 0, 
            dilation: _size_any_t = 1,
            return_indices: bool = False, 
            ceil_mode: bool = False
            )
        # kernel_size: 窗口尺寸；整数或者元组
        # stride: 窗口的步幅，默认就等于窗口尺寸；整数或者元组
        # padding: 数据hw方向上的零填充层数；整数或者元组
        # dilation: 一个控制窗口中元素步幅的参数；整数或者元组
        # return_indices: 如果等于True，会返回输出最大值的序号，对于上采样操作会有帮助
        # ceil_mode: 如果等于True，计算输出数据的hw的时候，会使用向上取整，代替默认的向下取整的操作
        """
        kernel_size = (2, 3)
        stride = (2, 1)
        self.pool = nn.MaxPool2d(kernel_size=kernel_size, stride=stride)

        """
        #-------------------------------线性连接层--------------------------------#
        对输入数据做线性变换(列变换)
        Linear(in_features: int, out_features: int, bias: bool = True)
        # in_features : N-每个输入样本的大小
        # out_features : M-每个输出样本的大小
        # bias : 若设置为False，这层不会学习偏置。默认值：True
        定义了一个线性层，它接收一个(K×N)大小的张量，输出一个(K×M)大小的张量
        y = xA.T + b
        权重矩阵放在PyTorch LinearLayer类中，是由PyTorch创建。
        PyTorch LinearLayer类使用传递给构造函数的数字 N 和 M 来创建一个(N×M)大小的权重矩阵。
        """
        in_features = 4
        out_features = 2
        self.fc = nn.Linear(in_features=in_features, out_features=out_features)

    def layer_basic_linear(self):
        """
        线性全连接层变换示例，示例中线性连接层的参数在init函数中指定，
        输入的X张量的数量则设为随机值
        """
        N = self.fc.in_features
        M = self.fc.out_features
        input0 = torch.randint(0, 10, (random.randint(1, 8), N), dtype=torch.float32)
        print("\033[1;31m示例中输入的张量X的维数为1-8之间的随机值\033[0m")
        print("以下是线性全连接层示例：")
        print("线性全连接层的输入有 {0}(N) 个特征，输出有 {1}(M) 个特征".format(N, M))
        print("Linear 层的工作原理是，使用随机设定的权重矩阵将一个 {0} 空间映射到一个 {1} 空间".format(N, M))
        print("输入张量X为")
        print(input0.data)
        print("输出张量Y为")
        output = self.fc(input0.data)
        print(output)
        print("————————————————————————————————————————————————————\n可以查看大小为({0}×{1})的列变换矩阵：".format(M, N))
        print(self.fc.weight.data)
        print("注意:\033[1;33m线性运算中实际为Y = X*A.T + b，所以A是({0}×{1})大小的\033[0m".format(M, N))
        print("以及线性变换的偏置矩阵:")
        print(self.fc.bias.data)
        print("————————————————————————————————————————————————————\n其实可以通过.weight.data显式地改变权重矩阵:")
        self.fc.weight.data = torch.eye(M, N).data
        print("看一看改变后的列变换矩阵:")
        print(self.fc.weight.data)
        print("和偏置矩阵")
        self.fc.bias.data = torch.zeros((1, M)).data
        print(self.fc.bias.data)

        print("再看看权重改变后的输出:")
        output = self.fc(input0.data)
        print(output)
        pass

    def layer_basic_Conv2d(self):
        """
        二维卷积层的使用示例
        :return:
        """
        k = self.conv.kernel_size
        C_in = self.conv.in_channels
        C_out = self.conv.out_channels
        t = random.randint(5, 9)
        h = random.randint(5, 9)
        w = random.randint(5, 9)
        input0 = torch.randint(10, (t, C_in, h, w), dtype=torch.float32)
        output = self.conv(input0)
        output_size = output.size()
        h1 = output_size[2]
        w1 = output_size[3]

        print("\033[1;31m-----------二维卷积层使用示例-------------\033[0m")
        print("官方中文教程请看 https://pytorch-cn.readthedocs.io/zh/latest/package_references/torch-nn/#_1")
        print("\033[1;32m示例中输入的张量X的维数为t={0}, h={1}, w={2}，t∈(5,9),h∈(5,9),w∈(5,9)\033[0m".format(t, w, h))
        print("卷积层的输入通道\033[1;31mC_in为 {0} \033[0m，输出通道\033[1;34mC_out为 {1} \033[0m，卷积核大小\033[1;33mk为 {2} \033[0m"
              .format(C_in, C_out, k))

        print("卷积层的输入必须和卷积层参数保持一致，\n卷积层输入张量X的大小为(N, \033[1;31mC_in\033[0m, H_in, W_in)")
        print("* 输入X的尺寸为：")
        print("* ({0}, \033[1;31m{1}\033[0m, {2}, {3})".format(t, C_in, h, w))
        print("卷积层的输出张量Y大小为(N, \033[1;34mC_out\033[0m,H_out,W_out),其中，H_out和W_out由原始尺寸和 k 来决定")
        print("* 输出Y的尺寸为：")
        print("* ({0}, \033[1;34m{1}\033[0m, \033[1;33m{2}\033[0m, \033[1;33m{3}\033[0m)".format(t, C_out, h1, w1))
        print("由于步长为默认值{0},卷积核大小k为{5},所以张量X的\n"
              "上下两侧各减去 (k[0] -1)/2= {1}，即h= h0- 2* {1}= \033[1;33m{3}\033[0m\n"
              "左右两侧各减去 (k[1] -1)/2= {2}，即w= w0- 2* {2}= \033[1;33m{4}\033[0m"
              .format(self.conv.stride, (k[0] - 1) / 2, (k[1] - 1) / 2, h1, w1, k))
        print("\n----------------------------------------------")
        print("* 可以查看卷积核的具体参数：(通过self.conv.weight.data调用)")
        print("卷积核的权重张量为:")
        print(self.conv.weight.data.size())
        pass

    def layer_basic_Pool(self):
        """
        二维最大池化层的使用示例
        :return:
        """
        print("\033[1;31m-----------二维最大池化层使用示例-------------\033[0m")
        print("官方中文教程请看 https://pytorch-cn.readthedocs.io/zh/latest/package_references/torch-nn/#_2")
        print("最大池化层的作用：\n"
              "（1）首要作用，下采样\n"
              "（2）降维、去除冗余信息、对特征进行压缩、简化网络复杂度、减小计算量、减小内存消耗等\n"
              "（3）实现非线性、\n"
              "（4）扩大感知野、\n"
              "（5）实现不变性，其中不变形性包括，平移不变性、旋转不变性和尺度不变性。\n")
        print("------------------------------------------------------")
        print("对于输入信号的输入通道，提供1维最大池化（max pooling）操作")
        print("输入的大小是(N,C,H,W)，那么输出的大小是(N,C,H_out,W_out)")
        N = 1
        C = 1
        H = random.randint(5, 8)
        W = random.randint(5, 8)
        print("\033[1;32m示例中输入的张量X的维数为N={0}, C={1}, H={2}，W={3}; N=1, C=1, H∈(5,8), W∈(5,8)\033[0m".format(N, C, H, W))

        k = self.pool.kernel_size
        s = self.pool.stride
        print("池化层的参数: 核大小={0}, 步长={1}".format(k, s))

        input0 = torch.randint(10, (N, C, H, W), dtype=torch.float32)
        print("输入张量的随机尺寸为\t{0}".format(input0.size()))
        print(input0.data)
        output = self.pool(input0)
        print("输出张量的尺寸为\t\t{0}".format(output.size()))
        print(output.data)
        print("\033[1;33m本质是移动大小为{0}的滑窗，提取滑窗中的最大值\033[0m".format(k))
        print("* 注意:\033[1;31m从左上角开始，末尾方向如果不能充满滑窗，则末尾数据会被丢弃，不作计算 \033[0m")

        pass



